package algotithm.leetcode.unionfind.test721;

import java.util.*;

/**
 * @author zhouyanxiang
 * @Date 2021-04-2021/4/9-14:08
 */
class Solution {
    public List<List<String>> accountsMerge(List<List<String>> accounts) {
        Map<String,Integer> emailToIndex = new HashMap<>();
        Map<String,String> emailToName = new HashMap<>();
        int emailCount = 0;
        for (List<String> account : accounts) {
            String name = account.get(0);
            int size = account.size();
            for (int i = 1; i < size; i++) {
                String email = account.get(i);
                if (!emailToIndex.containsKey(email)) {
                    emailToIndex.put(email,emailCount++);
                    emailToName.put(email,name);
                }
            }
        }
        UnionFind uf = new UnionFind(emailCount);
        for (List<String> account : accounts) {
            String email = account.get(1);
            int firstIndex = emailToIndex.get(email);
            int size = account.size();
            for (int i = 2; i < size; i++) {
                String nextEmail = account.get(i);
                int nextIndex = emailToIndex.get(nextEmail);
                uf.union(firstIndex,nextIndex);
            }
        }
        Map<Integer,List<String>> indexToEmail = new HashMap<>();
        for (String email : emailToIndex.keySet()) {
            int index = uf.find(emailToIndex.get(email));
            List<String> account = indexToEmail.getOrDefault(index,new ArrayList<>());
            account.add(email);
            indexToEmail.put(index,account);
        }
        List<List<String>> res = new ArrayList<>();
        for (List<String> emails : indexToEmail.values()) {
            Collections.sort(emails);
            String name = emailToName.get(emails.get(0));
            List<String> account = new ArrayList<>();
            account.add(name);
            account.addAll(emails);
            res.add(account);
        }
        return res;
    }
}

class UnionFind{

    public int[] p;

    public  UnionFind(int x) {
        p = new int[x+1];
        for (int i = 1; i <= x; i++) {
            p[i] = i;
        }
    }

    public void union(int x, int y) {
        p[find(x)] = find(y);
    }

    public int find(int x) {
        if(p[x] != x) {
            p[x] = find(p[x]);
        }
        return p[x];
    }

}
